package py;

import lombok.extern.slf4j.Slf4j;

import java.time.Duration;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;

@Slf4j
public class DelayedTaskExecutor {
    /**
     * 延时任务调度线程池
     */
    private final ScheduledExecutorService scheduler;

    /**
     * 实际执行任务线程池
     */
    private final ExecutorService worker;

    /**
     * 本地任务池，
     * 用于保存 key -> future 的映射，
     * 其中 key 是任务的唯一标识，
     * future 是任务的执行 Future 对象；
     */
    private final Map<String, ScheduledFuture<?>> taskPool = new ConcurrentHashMap<>();

    /**
     * 每个任务的最大执行时长
     */
    private final Duration taskTimeout;

    public DelayedTaskExecutor(Duration taskTimeout) {
        this(null, null, taskTimeout);
    }

    public DelayedTaskExecutor(ScheduledExecutorService scheduler, ExecutorService worker, Duration taskTimeout) {
        this.scheduler = scheduler == null ? Executors.newScheduledThreadPool(4) : scheduler;
        this.worker = worker == null ? new ThreadPoolExecutor(8, 16, 60, TimeUnit.SECONDS, new ArrayBlockingQueue<>(100)) : worker;
        this.taskTimeout = taskTimeout == null ? Duration.ofSeconds(10) : taskTimeout;
    }

    /**
     * 延时执行任务，并保证相同 key 的任务互斥，只执行最后一个
     * 示例：
     * <blockquote><pre>
     * DefaultTask task = new DefaultTask(UUID.randomUUID().toString(), "检测到有人进入，开启灯泡", () -> {
     *     // 下发MQTT消息给指定的设备
     *     // ...
     * });
     * delayedTaskExecutor.doAfter("the-end-device-key", task, Duration.ofMillis(50));
     * </pre></blockquote>
     *
     * @param key   任务的key值，相同key值的任务具有互斥性，也就是说在本地待执行的任务池中，
     *              相同key值的任务只执行最后一个，而前面未执行的相同key值的任务会被清除；
     *              此值一般是终端设备的唯一标识，
     *              例如：设备的MAC地址、设备的唯一ID等；
     * @param task  待执行的任务对象
     * @param delay 延时多长时间后执行
     */
    public void doAfter(String key, Task task, Duration delay) {
        Objects.requireNonNull(key);
        Objects.requireNonNull(task);
        Objects.requireNonNull(delay);

        // 如果已有任务，取消并移除
        ScheduledFuture<?> oldFuture = taskPool.remove(key);
        if (oldFuture != null) {
            oldFuture.cancel(false);
            log.info("Task replaced for key={}, old removed", key);
        }

        // 包装任务
        Runnable wrapped = () -> {
            // 任务真正开始执行
            log.info("Task starting: id={}, desc={}", task.getId(), task.getDesc());
            // 先删除映射，确保在任务执行完成前不会被新任务替换
            // 任务执行完成后，会自动从映射中移除
            taskPool.remove(key);
            try {
                Future<?> f = worker.submit(task::handle);
                try {
                    f.get(taskTimeout.toMillis(), TimeUnit.MILLISECONDS);
                } catch (TimeoutException e) {
                    // 超时取消任务
                    f.cancel(true);
                    log.warn("Task timeout and cancelled: id={}", task.getId());
                }
            } catch (Exception e) {
                // 任务执行异常
                log.error("Task execution failed: id={}, desc={}, error={}", task.getId(), task.getDesc(), e.getMessage(), e);
            }
        };

        // 提交延时任务
        ScheduledFuture<?> future = scheduler.schedule(wrapped, delay.toMillis(), TimeUnit.MILLISECONDS);
        taskPool.put(key, future);

        log.info("Task scheduled: id={}, desc={}, delay={}", task.getId(), task.getDesc(), delay);
    }

    public void shutdown() {
        scheduler.shutdownNow();
        worker.shutdownNow();
    }

}
